package com.hxy.springboot.langchain4jspringboot.controller;

import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.aot.AotServices;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

@Slf4j
@RestController
@RequestMapping("/ai")
public class AiController {
    @Autowired
    private QwenChatModel qwenChatModel;
    @Autowired
    private QwenStreamingChatModel qwenStreamingChatModel;
    //定义代理接口-正常
    interface Assistant{
        String chat(String message);
    }
    //定义代理接口-流式
    interface AssistantStream{
        TokenStream chat(String message);
    }

    /**
     * 通义千问普通回答
     *
     * @param message
     * @return
     */
    @GetMapping("/qwen")
    public String QwenModel(@RequestParam("message") String message) {
        String chat = qwenChatModel.chat(message);
        return chat;
    }

    /**
     * 通义千问流式回答
     *
     * @param message
     * @return
     */
    @GetMapping("/streamQwen")
    public String streamQwen(@RequestParam("message") String message) {
        qwenStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String s) {
                System.out.println(s);
            }

            @Override
            public void onCompleteResponse(ChatResponse chatResponse) {
                System.out.println(chatResponse);
            }

            @Override
            public void onError(Throwable throwable) {
                System.out.println(throwable);
            }
        });
        return "success";
    }

    /**
     * 通义千问记忆回答-普通回答
     */
    @GetMapping("/memoryQwen")
    public String memoryQwen(@RequestParam("message") String message) {
        MessageWindowChatMemory memory = MessageWindowChatMemory.withMaxMessages(10);
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatMemory(memory)
                .chatLanguageModel(qwenChatModel).build();
        String chat = assistant.chat(message);
        System.out.println(chat);
        return chat;
    }

    /**
     * 通义千问记忆回答-流式回答
     * 这部分和普通回答的差别在于使用了Flux来读取流的内容
     */
    @GetMapping("/memoryStreamQwen")
    public Flux<String> memoryStreamQwen(@RequestParam("message") String message) {
        MessageWindowChatMemory memory = MessageWindowChatMemory.withMaxMessages(10);
        AssistantStream assistantStream = AiServices.builder(AssistantStream.class)
                .chatMemory(memory)
                .streamingChatLanguageModel(qwenStreamingChatModel).build();
        TokenStream chat = assistantStream.chat(message);
        Flux<String> flux = Flux.create(fluxSink -> {
            chat.onPartialResponse(s -> fluxSink.next(s))
                    .onCompleteResponse(chatResponse -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });
        return flux;
    }
}
